#include "tcp.h"
#include "eth.h"
#include "ip.h"
#include "posix_api.h"
#include "arp.h"
#include "util.h"



struct tcp_stream * tcp_stream_search(uint32_t s_ip, uint32_t d_ip, uint16_t s_port, uint16_t d_port){
    struct tcp_table * table = get_tcp_instance();
    struct tcp_stream * it;
    for(it = table->links; it != NULL; it = it->next){  // established 通过五元组查找
        if(it->s_ip == s_ip && it->d_ip == d_ip && it->s_port == s_port && it->d_port == d_port){
            return it;
        }
    }

    for(it = table->links; it != NULL; it = it->next){
        if(it->d_port == d_port && it->status == TCP_STATUS_LISTEN){   // listen 此时没有 源ip和端口，只能通过目的端口查找
            return it;
        }
    }

    return NULL;
}

struct tcp_stream * tcp_stream_create(uint32_t s_ip, uint32_t d_ip, uint16_t s_port, uint16_t d_port){
    struct tcp_stream * stream = rte_malloc("create stream", sizeof(struct tcp_stream), 0);
    if(stream == NULL){
        return  NULL;
    }
    stream->s_port = s_port;
    stream->d_port = d_port;
    stream->s_ip = s_ip;
    stream->d_ip = d_ip;
    stream->protocol = IPPROTO_TCP;

    stream->fd = -1;   // 新创建  不分配fd 这个在accept分配

    stream->status = TCP_STATUS_LISTEN;
    stream->rcv_buf = rte_ring_create("rec buf", RING_SIZE, rte_socket_id(), 0);
    stream->send_buf = rte_ring_create("send buf", RING_SIZE, rte_socket_id(), 0);

    uint32_t next_seed = time(NULL);
    stream->seq_num = rand_r(&next_seed) % TCP_MAX_SEQ;
    rte_memcpy(stream->local_mac, g_src_mac, RTE_ETHER_ADDR_LEN);

    pthread_cond_t blank_cond = PTHREAD_COND_INITIALIZER;
    rte_memcpy(&stream->cond, &blank_cond, sizeof(pthread_cond_t));
    pthread_mutex_t blank_mutex = PTHREAD_MUTEX_INITIALIZER;
    rte_memcpy(&stream->mtx, &blank_mutex, sizeof(pthread_mutex_t));

    return stream;
}

static void _encode_tcp_hdr_pkt(uint8_t * msg, uint16_t s_port, uint16_t d_port, uint32_t tx_seq_num, uint32_t rx_ack_seq_num,
    uint8_t data_offset, uint8_t  tcp_flags, uint16_t rx_win, uint8_t * data, int opt_len, uint32_t length){
    struct rte_ipv4_hdr * iphdr = (struct rte_ipv4_hdr *)(msg + sizeof(struct rte_ether_hdr));
    struct rte_tcp_hdr * tcphdr = (struct rte_tcp_hdr *)(iphdr + 1);
    tcphdr->src_port = s_port;
    tcphdr->dst_port = d_port;
    tcphdr->sent_seq = tx_seq_num;   // 序号 随机值 (字节序号)
    tcphdr->recv_ack = rx_ack_seq_num;  // 确认号
    tcphdr->data_off = data_offset;
    tcphdr->tcp_flags = tcp_flags;
    tcphdr->rx_win = rx_win;
    if(data != NULL){
        uint8_t * payload = (uint8_t*)(tcphdr + 1) + opt_len * sizeof(uint32_t);
        rte_memcpy(payload, data, length);
    }
    tcphdr->cksum = 0;
    tcphdr->cksum = rte_ipv4_udptcp_cksum(iphdr, tcphdr);
}


static struct rte_mbuf * _encode_tcp(uint8_t * msg, uint32_t src_ip, uint32_t dst_ip,
    uint8_t * src_mac, uint8_t * dst_mac, struct tcp_fragment * fragment, const unsigned total_len){
    encode_ether_hdr_pkt(msg, dst_mac, src_mac);
    encode_ip_hdr_pkt(msg, total_len - sizeof(struct rte_ether_hdr), IPPROTO_TCP, src_ip, dst_ip);
    _encode_tcp_hdr_pkt(msg, fragment->s_port, fragment->d_port, htonl(fragment->seq_num), htonl(fragment->ack_num)
        ,fragment->hdrlen_off, fragment->tcp_flags, fragment->windows, fragment->data, fragment->optlen, fragment->length);

    printf("__encode tcp:  data: %s,  length: %d  \n", fragment->data, fragment->length);
    return 0;
}

static struct rte_mbuf * _send_tcp(struct rte_mempool * pool, uint32_t src_ip, uint32_t dst_ip, uint8_t * src_mac, uint8_t * dst_mac, struct tcp_fragment * fragment){
    const unsigned total_len =fragment->length + sizeof(struct rte_ether_hdr) +
                    sizeof(struct rte_ipv4_hdr) + sizeof(struct rte_tcp_hdr) + fragment->optlen * sizeof(uint32_t);

    printf("----__send_tcp------total_len: %d-----------\n", total_len);
    struct rte_mbuf * mbuf = rte_pktmbuf_alloc(pool);
    if(!mbuf){
    error_print(" __send_tcp alloc mbuf");
    }
    mbuf->pkt_len = total_len;
    mbuf->data_len = total_len;
    uint8_t * pkt_data = rte_pktmbuf_mtod(mbuf, uint8_t *);
    _encode_tcp(pkt_data, src_ip, dst_ip, src_mac, dst_mac, fragment, total_len);
    return mbuf;
}

int tcp_pkg_process(struct rte_mbuf * m_buf){

}



int tcp_handle_listen(struct tcp_stream * stream, struct rte_tcp_hdr * tcphdr, struct rte_ipv4_hdr * iphdr){
    /**
     *  1. 创建一个半链接节点， fd没有分配
     *  2. 返回 ack 和 syn 包
     * */
    if(tcphdr->tcp_flags & RTE_TCP_SYN_FLAG){   // tcp 头的标志位 syn 是否置 1
        if(stream->status == TCP_STATUS_LISTEN){   // 避免重复发
            struct tcp_table * table = get_tcp_instance();
            // 无需考虑有重复的数据包 因为再次查找会查找到已经创建的链接
            struct tcp_stream * syn = tcp_stream_create(iphdr->src_addr, iphdr->dst_addr, tcphdr->src_port, tcphdr->dst_port);
            LL_ADD(syn, table->links);

            struct tcp_fragment * fragment = rte_malloc("h listen fragment", sizeof(struct tcp_fragment), 0);
            if(fragment == NULL)
                return -1;
            memset(fragment, 0, sizeof(struct tcp_fragment));

            fragment->s_port = tcphdr->dst_port;
            fragment->d_port = tcphdr->src_port;

            print_address("tcp   src: ", syn->s_ip);
            printf("----sport----%d \n", tcphdr->src_port);
            print_address("tcp   dst: ", syn->d_ip);
            printf("---dport-----%d \n", tcphdr->dst_port);

            fragment->seq_num = syn->seq_num;
            fragment->ack_num = ntohl(tcphdr->sent_seq) + 1;
            syn->ack_num = fragment->ack_num;

            fragment->tcp_flags = (RTE_TCP_SYN_FLAG | RTE_TCP_ACK_FLAG);
            fragment->windows = TCP_INITIAL_WINDOW;
            fragment->hdrlen_off = 0x50;        // 首部长度20字节

            fragment->data = NULL;
            fragment->length = 0;
            rte_ring_mp_enqueue(syn->send_buf, fragment);
            syn->status = TCP_STATUS_SYN_RCVD;
        }
    }
    return 0;
}

int tcp_handle_syn_rcvd(struct tcp_stream * stream, struct rte_tcp_hdr * tcphdr){
    if(tcphdr->tcp_flags & RTE_TCP_ACK_FLAG){
        if(stream->status == TCP_STATUS_SYN_RCVD){
            uint32_t ack_num = ntohl(tcphdr->recv_ack);
            if(ack_num == stream->seq_num + 1){
                //
            }
            stream->status = TCP_STATUS_ESTABLISHED;

            struct tcp_stream * listener = tcp_stream_search(0,0,0, stream->d_port);   // 查找处于 listen 状态的链接 也就是listenfd
            if(listener == NULL){
                error_print("tcp handle syn recvd listener");
            }
            pthread_mutex_lock(&listener->mtx);
            pthread_cond_signal(&listener->cond);   //  唤醒 accept
            pthread_mutex_unlock(&listener->mtx);

#if ENABLE_EPOLL
            // 第一个调入点
            struct tcp_table * table = get_tcp_instance();
            epoll_event_cb(table->ep,listener->fd, EPOLLIN);
#endif            
        }
    }
    return 0;
}

int tcp_enqueue_recv_buf(struct tcp_stream * stream, struct rte_tcp_hdr * tcphdr, int tcp_len){
    printf("-----------tcp_enqueue_recv_buf---------------->> \n");
    struct tcp_fragment * recv_fragment = rte_malloc("tcp fragment", sizeof(struct tcp_fragment), 0);
    if(recv_fragment == NULL){
        return -1;
    }
    memset(recv_fragment, 0, sizeof(struct tcp_fragment));
    recv_fragment->d_port = ntohs(tcphdr->dst_port);
    recv_fragment->s_port = ntohs(tcphdr->src_port);

    uint8_t hdrlen = tcphdr->data_off >> 4;
    int payload_len = tcp_len - hdrlen * 4;
    if(payload_len > 0){
        uint8_t * payload = (uint8_t *)tcphdr + hdrlen * 4;  // tcp 数据位置
        recv_fragment->data = rte_malloc("recv fragment data", payload_len + 1, 0);
        if(recv_fragment->data == NULL){
            rte_free(recv_fragment);
            return -1;
        }
        memset(recv_fragment->data, 0, payload_len + 1);   // 末尾 \0
        rte_memcpy(recv_fragment->data, payload, payload_len);
        recv_fragment->length = payload_len;
    }else if(payload_len == 0){
        recv_fragment->length = 0;
        recv_fragment->data = NULL;
    }
    rte_ring_mp_enqueue(stream->rcv_buf, recv_fragment);
    pthread_mutex_lock(&stream->mtx);
    pthread_cond_signal(&stream->cond);   // 唤醒 recv 阻塞
    pthread_mutex_unlock(&stream->mtx);

    return 0;
}

int tcp_send_ack(struct tcp_stream * stream, struct rte_tcp_hdr * tcphdr){
    struct tcp_fragment * ack_fragment = rte_malloc("ack fragment", sizeof(struct tcp_fragment), 0);
    if(ack_fragment == NULL){
        return -1;
    }
    memset(ack_fragment, 0, sizeof(struct tcp_fragment));

    ack_fragment->d_port = tcphdr->src_port;
    ack_fragment->s_port = tcphdr->dst_port;

    printf("ng_tcp_send_ackpkt: %d, %d\n", stream->seq_num, ntohs(tcphdr->sent_seq));

    ack_fragment->ack_num = stream->ack_num;
    ack_fragment->seq_num = stream->seq_num;

    ack_fragment->tcp_flags = RTE_TCP_ACK_FLAG;
    ack_fragment->windows = TCP_INITIAL_WINDOW;
    ack_fragment->hdrlen_off = 0x50;
    ack_fragment->data = NULL;
    ack_fragment->length = 0;

    rte_ring_mp_enqueue(stream->send_buf, ack_fragment);
    return 0;
}

int tcp_handle_established(struct tcp_stream * stream, struct rte_tcp_hdr * tcphdr, int tcp_len){
    printf("------------tcp_handle_established--------------->> \n");

    /**
     *  收发数据
     * */
    if(tcphdr->tcp_flags & RTE_TCP_SYN_FLAG){
        printf("------------RTE_TCP_SYN_FLAG--------------->> \n");
    }else if(tcphdr->tcp_flags & RTE_TCP_PSH_FLAG){  // 不需要等后面的数据包，这个包现在立刻马上处理
        printf("-----------------RTE_TCP_PSH_FLAG---------->> \n");
        tcp_enqueue_recv_buf(stream, tcphdr, tcp_len);  // 收到数据包 放到 recv_buf 回应ack

#if ENABLE_EPOLL
        // 第二次调入点
        struct tcp_table * table = get_tcp_instance();
        epoll_event_cb(table->ep,stream->fd, EPOLLIN);
#endif

        uint8_t hdrlen = tcphdr->data_off >> 4;
        int payloadlen = tcp_len - hdrlen * 4;

        stream->ack_num = stream->ack_num + payloadlen;
        stream->seq_num = ntohl(tcphdr->recv_ack);

        tcp_send_ack(stream, tcphdr);
    }else if(tcphdr->tcp_flags & RTE_TCP_ACK_FLAG|RTE_TCP_FIN_FLAG){ // 已经是established 收到 ack， 不会引起状态迁移
        if(tcphdr->tcp_flags & RTE_TCP_FIN_FLAG){
            printf("-----------RTE_TCP_FIN_FLAG---------------->> \n");
            stream->status = TCP_STATUS_CLOSE_WAIT;
            tcp_enqueue_recv_buf(stream, tcphdr, tcphdr->data_off >> 4);

#if ENABLE_EPOLL
            // 第三个调用点
            struct tcp_table * table = get_tcp_instance();
            epoll_event_cb(table->ep,stream->fd, EPOLLIN);
#endif

            stream->ack_num = stream->ack_num + 1;
            stream->seq_num = ntohl(tcphdr->recv_ack);
            tcp_send_ack(stream, tcphdr);
        }else{
            printf("-------------RTE_TCP_ACK_FLAG---dd----------->> \n");
        }
    }
    return 0;
}

int tcp_handle_close_wait(struct tcp_stream * stream, struct rte_tcp_hdr * tcphdr){
    printf("-------------tcp_handle_close_wait-------------->> \n");

    if (tcphdr->tcp_flags & RTE_TCP_FIN_FLAG) { //
        if (stream->status == TCP_STATUS_CLOSE_WAIT) {

        }
    }
    return 0;
}

int tcp_handle_last_ack(struct tcp_stream * stream, struct rte_tcp_hdr * tcphdr){
    printf("-------------tcp_handle_last_ack-------------->> \n");

    if(tcphdr->tcp_flags & RTE_TCP_ACK_FLAG){
        if(stream->status == TCP_STATUS_LAST_ACK){
            stream->status = TCP_STATUS_CLOSED;
            printf("last ack \n");
            struct tcp_table * table = get_tcp_instance();
            LL_REMOVE(stream, table->links);
            rte_ring_free(stream->send_buf);
            rte_ring_free(stream->rcv_buf);
            rte_free(stream);
        }
    }
}

int tcp_process(struct rte_mbuf * m_buf){  // 解析 tcp 数据包 根据不同的状态进行不同的处理
    struct rte_ipv4_hdr * iphdr = rte_pktmbuf_mtod_offset(m_buf, struct rte_ipv4_hdr *, sizeof(struct rte_ether_hdr));
    struct rte_tcp_hdr * tcphdr = (struct rte_tcp_hdr *)(iphdr + 1);

    uint16_t tcp_ck_sum = tcphdr->cksum;
    tcphdr->cksum = 0;
    uint16_t ck_sum = rte_ipv4_udptcp_cksum(iphdr, tcphdr);  // 计算后是本地值，需要将网络字节转换成本地

    if(ck_sum != tcp_ck_sum){   //  本地计算的cksum  和 网络传过来的cksum 对比
        printf("ck sum: %x,  tcp ck_sum :  %x \n", ck_sum, tcp_ck_sum);
        return -1;
    }
    // 查找tcb
    struct tcp_stream * stream = tcp_stream_search(iphdr->src_addr, iphdr->dst_addr, tcphdr->src_port, tcphdr->dst_port);
    if(stream == NULL){
        return -2;
    }

    printf("--------stream-status---%d--------------->> \n", stream->status);


    switch (stream->status) {
        /* client */
        case TCP_STATUS_CLOSED:
            break;

        /* server */
        case TCP_STATUS_LISTEN:   // SYN_RCVD 发 RST
            tcp_handle_listen(stream, tcphdr, iphdr);
            break;

        /* client */
        case TCP_STATUS_SYN_SENT:  // LISTEN  发 syn；  通常不会使用一个listen_fd 调用connect连接远程服务器
            break;

        /* server */
        case TCP_STATUS_SYN_RCVD:  // 1. LISTEN  收到 syn 发送syn 和 ack；  2. SYN_SENT 收到 syn 发送 syn ack
            tcp_handle_syn_rcvd(stream, tcphdr);
            break;

        /* server client */
        case TCP_STATUS_ESTABLISHED:  // 1. SYN_RCVD 收到 ack；  2. SYN_SENT 收到 syn ack 发送 ack
        {
            // 打印一个tcp 出来
            uint8_t hdrlen = tcphdr->data_off & 0xF0;    // tcp 头长度  &  0xF0  ---得到---> 0x50
            hdrlen >>= 4;                                // 得到 首部长度
            uint8_t * payload = (uint8_t *)tcphdr + hdrlen * 4;  // 偏移完头部长度 打印数据
            printf("payload:  %s \n", payload);

            int tcp_len = ntohs(iphdr->total_length) - sizeof(struct rte_ipv4_hdr);
            tcp_handle_established(stream, tcphdr, tcp_len);
            break;
        }

        //  主动关闭
        case TCP_STATUS_FIN_WAIT_1:  //  ESTABLISHED 或 SYN_RCVD ---> 进程关闭，发送 fin
            break;
        case TCP_STATUS_FIN_WAIT_2:  //  FIN_WAIT_1  收到ack
            break;
        case TCP_STATUS_CLOSING:     //  FIN_WAIT_1  收到 fin 发送 ack
            break;
        case TCP_STATUS_TIME_WAIT:  // 等待超时 2MSL
                                    //  1. FIN_WAIT_2  收到 fin 发送 ack；  2. FIN_WAIT_1 收到 ack 发送 ack； 3. CLOSING 收到 ack
            break;

        //  被动关闭
        case TCP_STATUS_CLOSE_WAIT:  // ESTABLISHED 收到 fin 发送 ack
            tcp_handle_close_wait(stream, tcphdr);
            break;
        case TCP_STATUS_LAST_ACK:    // CLOSE_WAIT  关闭应用程序，发送 fin    收到 ack 关闭；
            tcp_handle_last_ack(stream, tcphdr);
            break;
        default:
            break;
    }
    return 0;
}

int tcp_out(struct rte_mempool * pool){  // 发送 tcp 数据包【放到 out ring】
    struct tcp_table * table = get_tcp_instance();
    struct tcp_stream * stream;
    for (stream = table->links; stream != NULL; stream = stream->next){
        if(stream->send_buf == NULL){  // listener 阶段
            continue;
        }
        struct tcp_fragment * fragment = NULL;
        int number_send = rte_ring_mc_dequeue(stream->send_buf, (void **)&fragment); // 取出多个包
        if(number_send < 0){  // 没有需要发送的tcp包
            continue;
        }

        uint8_t * dst_mac = get_dst_mac_address(stream->s_ip);
        if(dst_mac == NULL){
            // 发送arp (struct rte_mempool * m_pool, uint16_t opcode, uint8_t * dst_mac, uint32_t s_ip, uint32_t d_ip)
            struct rte_mbuf * arp_buf = arp_send(pool, RTE_ARP_OP_REQUEST, g_default_arp_mac, stream->d_ip, stream->s_ip);
            struct inout_ring * ring = ring_instance();
            rte_ring_mp_enqueue_burst(ring->out, (void **)&arp_buf, 1, NULL);
            rte_ring_mp_enqueue(stream->send_buf, fragment);
        }else{
            struct rte_mbuf * tcp_buf = _send_tcp(pool, stream->d_ip, stream->s_ip, stream->local_mac, dst_mac, fragment);
            struct inout_ring * ring = ring_instance();
            rte_ring_mp_enqueue_burst(ring->out, (void **)&tcp_buf, 1, NULL);
            if(fragment->data != NULL){
                rte_free(fragment->data);
            }
            rte_free(fragment);
        }
    }
    return 0;
}

int tcp_server_entry(__attribute__((unused)) void * arg){
    printf("----------tcp_server_entry----------------->> \n");
    int listen_fd = m_socket(AF_INET, SOCK_STREAM, 0);
    if(listen_fd == -1){
        return -1;
    }
    struct sockaddr_in server_address;
    bzero(&server_address, sizeof(struct sockaddr_in));
    server_address.sin_family = AF_INET;
    server_address.sin_port = htons(G_PORT);
    server_address.sin_addr.s_addr = htonl(INADDR_ANY);
    m_bind(listen_fd, (struct sockaddr *)&server_address, sizeof(server_address));
    m_listen(listen_fd, 64);

#if ENABLE_EPOLL
    int epfd = nepoll_create(1);  // 创建了一个 event_poll结构体
    struct epoll_event ev;
    ev.events = EPOLLIN;
    ev.data.fd = listen_fd;
    nepoll_ctl(epfd, EPOLL_CTL_ADD, listen_fd, &ev);
#endif

    while(1){
        struct sockaddr_in client;
        socklen_t len = sizeof(client);
        int conn_fd = m_accept(listen_fd, (struct sockaddr *)&client, &len);  // 阻塞等待 client连接
        printf("---------conn_fd------%d------------>> \n", conn_fd);
        char buf[TCP_APP_RECV_BUFFER_SIZE] = {0};
        while(1){
            
#if ENABLE_EPOLL
            int nready = nepoll_wait(epfd, events, 128, -1);
#else
            int nready = m_recv(conn_fd, buf, TCP_APP_RECV_BUFFER_SIZE, 0);  // 从conn_fd recv_buf 拿节点 如果没有会阻塞
#endif
            if(nready > 0){
                printf("recv:  %s \n", buf);
                m_send(conn_fd, buf, nready, 0); // 回显
            }else if(nready == 0){
                printf("---------444------------------>> \n");
                m_close(conn_fd);
                break;
            }else{  // nonblock
                printf("----------555----------------->> \n");
            }
        }
    }
    m_close(listen_fd);
}
